import aiohttp
import asyncio
import json
import time
import pandas as pd
from datetime import datetime
from itertools import islice

async def calculate_loss(session, base_url, auth_token, loss_calc_payload, transactions):
    """异步股票盈亏试算完整业务流程"""
    headers = {
        'platform': 'ios',
        'appversion': '1.0.0',
        'language': 'zh-CN',
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {auth_token}'
    }

    try:
        # 创建股票试算记录
        loss_calc_url = f"{base_url}/v1/loss_caculation/add_update"
        async with session.post(loss_calc_url, headers=headers, json=loss_calc_payload) as response:
            result = await response.json()
            loss_id = result.get('data', {}).get('id')
            if not loss_id:
                return None, None

        # 提交股票交易记录
        detail_url = f"{base_url}/v1/loss_caculation_detail/add_update"
        for transaction in transactions:
            detail_payload = {
                "lossId": loss_id,
                "operationTime": transaction['operationTime'],
                "direction": 1 if transaction['direction'] == 'buy' else 0,
                "price": transaction['price'],
                "qty": transaction['qty'],
                "amount": transaction['amount']
            }
            async with session.post(detail_url, headers=headers, json=detail_payload) as response:
                await response.json()

        # 执行盈亏试算
        calculate_url = f"{base_url}/v1/loss_caculation/begin_caculat"
        params = {'caculateId': loss_id}
        async with session.get(calculate_url, headers=headers, params=params) as response:
            await response.json()
        
        return loss_calc_payload['stockNo'], loss_id

    except Exception as e:
        print(f"处理股票 {loss_calc_payload['stockNo']} 失败: {str(e)}")
        return None, None

async def process_stock_batch(session, stock_batch, auth_token):
    """异步处理一批股票"""
    tasks = []
    for stock_code, group in stock_batch:
        stock_name = group['stockName'].iloc[0]
        
        # 准备loss_calc_payload
        loss_calc_payload = {
            "stockNo": str(stock_code).zfill(6),
            "stockName": stock_name,
            "brokerName": "中信证券",
            "stockBroker": "600030"
        }
        
        # 准备transactions
        transactions = []
        for _, row in group.iterrows():
            transaction = {
                "operationTime": int(time.mktime(datetime.strptime(row['成交日期'], "%Y-%m-%d").timetuple())),
                "direction": row['买卖方向'].lower(),
                "price": float(row['成交价格']),
                "qty": int(row['成交数量']),
                "amount": float(row['成交额'])
            }
            transactions.append(transaction)
        
        # 创建异步任务
        task = calculate_loss(session, "http://119.91.232.205:6666", auth_token, 
                            loss_calc_payload, transactions)
        tasks.append(task)
    
    # 并发执行所有任务
    results = await asyncio.gather(*tasks)
    
    # 处理结果
    processed_results = []
    for stock_code, loss_id in results:
        if stock_code and loss_id:
            stock_name = next(group['stockName'].iloc[0] for code, group in stock_batch if str(code).zfill(6) == stock_code)
            processed_results.append({
                'stock_code': stock_code,
                'stock_name': stock_name,
                'loss_id': loss_id
            })
            print(f"股票 {stock_code} {stock_name} 试算成功，loss_id: {loss_id}")
    
    return processed_results

async def main():
    try:
        # 登录获取token
        base_url = "http://119.91.232.205:6666"
        login_url = f"{base_url}/v1/user/login"
        login_payload = {
            "email": "test@test.com",
            "password": "Aa123456",
            "keep": True
        }
        headers = {
            'platform': 'ios',
            'appversion': '1.0.0',
            'language': 'zh-CN',
            'Content-Type': 'application/json'
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.post(login_url, headers=headers, json=login_payload) as response:
                result = await response.json()
                auth_token = result.get('data', {}).get('token')
        
            if not auth_token:
                print("登录失败")
                return

            # 读取Excel文件
            df = pd.read_excel("./test/1000_stocks_test_data.xlsx")
            df.columns = ['stockNo', 'stockName', '成交日期', '买卖方向', '成交数量', '成交价格', '成交额']
            
            # 按股票代码分组
            grouped = df.groupby('stockNo')
            
            # 存储所有loss_id
            all_loss_ids = []
            
            # 每批处理100个股票
            batch_size = 100
            
            # 将分组数据分批处理
            it = iter(grouped)
            while True:
                # 获取下一批股票
                batch = list(islice(it, batch_size))
                if not batch:
                    break
                
                # 处理当前批次
                results = await process_stock_batch(session, batch, auth_token)
                all_loss_ids.extend(results)
                
                # 等待1秒后处理下一批
                # await asyncio.sleep(1)
        
        # 保存所有loss_id到文件
        with open('./test/test_008_loss_ids_1000.json', 'w', encoding='utf-8') as f:
            json.dump(all_loss_ids, f, ensure_ascii=False, indent=2)
        
        print(f"\n测试完成:")
        print(f"成功处理的股票数量: {len(all_loss_ids)}")
        print(f"Loss ID已保存到 test_008_loss_ids_1000.json")
        
    except Exception as e:
        print(f"程序执行出错: {str(e)}")

if __name__ == "__main__":
    asyncio.run(main())



    """

    #TODO 并发100个请求，计算1000个股票，任务卡死

    status count(1)
    1	27
    4	706
    5	265
    10	2
    

    SELECT
	`status`,
	COUNT(1) AS count 
FROM
	stock_loss_caculation
WHERE
	created_at BETWEEN '1735410177' 
	AND '1735410326' 
GROUP BY
	`status` 
ORDER BY
	`status`;
    
    """